package connect

import (
	"fmt"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"net"
	"time"
)

type Host struct {
	IP         string
	SSHPort    int
	Username   string
	Password   string
	SSHKey     string
	Passphrase string        // ssh key 的密码
	Timeout    time.Duration // 超时时间
	SSHClient  *ssh.Client
	SftpClient *sftp.Client
	JumpServer *JumpServer
}

func NewHost(ip, username, password, sshKey, passphrase string, sshPort int, jumpServer *JumpServer, timeout time.Duration) *Host {
	return &Host{
		IP:         ip,
		SSHKey:     sshKey,
		Username:   username,
		Password:   password,
		SSHPort:    sshPort,
		Passphrase: passphrase,
		Timeout:    timeout,
		JumpServer: jumpServer,
	}
}

func (h *Host) Open() error {
	if h.JumpServer == nil {
		return h.openWithoutJumpServer()
	}
	return h.openWithJumpServer()
}

func (h *Host) openWithoutJumpServer() error {
	auth, err := sshClientConfig(h.Password, h.SSHKey, h.Passphrase) // 生成密钥
	if err != nil {
		return err
	}
	config := &ssh.ClientConfig{
		User: h.Username,
		Auth: auth,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: h.Timeout,
	}
	h.SSHClient, err = ssh.Dial("tcp", fmt.Sprintf("%s:%d", h.IP, h.SSHPort), config)
	return err
}

func (h *Host) openWithJumpServer() error {
	conn, err := h.JumpServer.dial(h.IP, h.SSHPort)
	if err != nil {
		return err
	}
	auth, err := sshClientConfig(h.Password, h.SSHKey, h.Passphrase)
	if err != nil {
		return err
	}
	// 生成目标机器的ssh client配置
	config := &ssh.ClientConfig{
		User: h.Username,
		Auth: auth,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: h.Timeout,
	}

	clientConn, channels, requests, err := ssh.NewClientConn(conn, fmt.Sprintf("%s:%d", h.IP, h.SSHPort), config)
	if err != nil {
		return err
	}
	h.SSHClient = ssh.NewClient(clientConn, channels, requests)
	return nil
}

func (h *Host) OpenSftp() (err error) {
	if h.JumpServer == nil {
		if err = h.openWithoutJumpServer(); err != nil {
			return err
		}
	} else {
		if err = h.openWithJumpServer(); err != nil {
			return err
		}
	}
	h.SftpClient, err = sftp.NewClient(h.SSHClient)
	return err
}

func (h *Host) Close() {
	if h.SftpClient != nil {
		_ = h.SftpClient.Close()
	}
	_ = h.SSHClient.Close()
}
